import torch
import torchvision
from torch import nn
from torch.nn import functional as F
from torch.utils.data import DataLoader

class LeNet5(nn.Module):
    def __init__(self):
        super(LeNet5, self).__init__()
        self.conv_unit = nn.Sequential(
            nn.Conv2d(3, 6, kernel_size=5, stride=1, padding=0),
            nn.AvgPool2d(kernel_size=2, stride=2, padding=0),
            nn.Conv2d(6, 16, kernel_size=5, stride=1, padding=0),
            nn.AvgPool2d(kernel_size=2, stride=2, padding=0),

        )
        self.fc_unit = nn.Sequential(
            nn.Linear(16*5*5, 120),
            nn.ReLU(),
            nn.Linear(120, 84),
            nn.ReLU(),
            nn.Linear(84, 10)
        )

    def forward(self, x):
        batch_size = x.size(0)
        x = self.conv_unit(x)
        x = x.view(batch_size, 16*5*5)
        x = self.fc_unit(x)
        return x





train_data = torchvision.datasets.CIFAR10(root="./dataset", train=True, transform=torchvision.transforms.Compose([
    #torchvision.transforms.Resize(32, 32),
    torchvision.transforms.ToTensor()
]), download=True)

train_dataloader = DataLoader(train_data, batch_size=32, shuffle=True)


test_data = torchvision.datasets.CIFAR10(root="./dataset", train=False, transform=torchvision.transforms.Compose([
    #torchvision.transforms.Resize(32, 32),
    torchvision.transforms.ToTensor()
]), download=True)
test_dataloader = DataLoader(test_data, batch_size=32, shuffle=True)



model = LeNet5()
criterion = nn.CrossEntropyLoss()
optimzer = torch.optim.Adam(model.parameters(), lr=1e-3)


for epoch in range(100):
    model.train()
    for batchidx, (x, label) in enumerate(train_dataloader):
        x = model(x)
        loss = criterion(x, label)

        optimzer.zero_grad()
        loss.backward()
        optimzer.step()
    print(epoch, loss.item())

    model.eval()
    with torch.no_grad(): #表示下面的计算不需要构建计算图
        total_correct = 0
        total_num = 0
        for x, label in test_dataloader:
            x = model(x)
            pred = x.argmax(dim=1)
            total_correct += torch.eq(pred, label).float().sum().item()
            total_num += x.size(0)
        acc = total_correct / total_num
        print(epoch, acc)